from agent.lmstudio_agent import LLMAgent

# 配置选项
USE_LOCAL_MODEL = True  # 设置为False可切换到在线大模型

# 在线大模型配置 (仅当USE_LOCAL_MODEL=False时有效)
ONLINE_CONFIG = {
    "api_key": "your_api_key_here",  # 替换为实际的API密钥
    "provider": "openai",
    "default_model": "gpt-3.5-turbo"
}

# 本地模型配置 (仅当USE_LOCAL_MODEL=True时有效)
LOCAL_CONFIG = {
    "base_url": "http://localhost:1234/v1",
    "default_model": None  # 使用LM Studio默认模型
}

def main():
    # 初始化Agent
    if USE_LOCAL_MODEL:
        print("\n===== 使用本地LM Studio模型 =====")
        try:
            agent = LLMAgent(
                model_type="local",
                base_url=LOCAL_CONFIG["base_url"],
                default_model=LOCAL_CONFIG["default_model"]
            )
        except Exception as e:
            print(f"初始化本地模型失败: {e}")
            return
    else:
        print("\n===== 使用在线大模型 =====")
        try:
            agent = LLMAgent(
                model_type="online",
                api_key=ONLINE_CONFIG["api_key"],
                online_provider=ONLINE_CONFIG["provider"],
                default_model=ONLINE_CONFIG["default_model"]
            )
        except Exception as e:
            print(f"初始化在线模型失败: {e}")
            print("提示: 请检查API密钥是否有效，并确保网络连接正常。")
            return

    # 列出可用模型
    print("\n===== 可用模型 =====")
    models = agent.list_models()
    print(models)

    # 简单文本生成示例
    print("\n===== 文本生成示例 =====")
    prompt = "请解释什么是人工智能"
    response = agent.generate_text(prompt)
    print(f"提示: {prompt}")
    print(f"响应: {response}")

    # 多轮对话示例
    print("\n===== 多轮对话示例 =====")
    messages = [
        {"role": "user", "content": "你好，你是谁？"},
        {"role": "assistant", "content": "我是一个AI助手。"},
        {"role": "user", "content": "你能告诉我什么是机器学习吗？"}
    ]
    response = agent.chat_completion(messages)
    if "choices" in response and len(response["choices"]) > 0:
        print(f"AI助手: {response['choices'][0]['message']['content']}")
    else:
        print("没有获取到响应")


if __name__ == "__main__":
    main()
